# 3D 高楼人群疏散模拟：基于元胞自动机与动态出口选择策略
# 主要依赖库：numpy, matplotlib (3D 动画), random

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random

# 模拟参数配置
BUILDING_SIZE = (3, 10, 10)  # (楼层数, 行, 列)
POPULATION_DENSITY = 0.2     # 每层密度
STAIR_POSITIONS = [(5, 5)]   # 所有楼层楼梯在相同位置
EXITS = [(0, 0, 0), (0, 9, 9)]  # 出口位于最底层的两个角落
MAX_ENERGY = 100

# 体力消耗系数
ALPHA = 1    # 水平移动消耗
BETA = 3     # 垂直移动消耗（楼梯）
GAMMA = 5    # 拥堵区额外惩罚

# 定义人群个体类
class Person:
    def __init__(self, pos):
        self.pos = pos
        self.energy = MAX_ENERGY
        self.finished = False

    def move(self, new_pos, energy_cost):
        self.pos = new_pos
        self.energy -= energy_cost
        if self.pos in EXITS:
            self.finished = True

# 初始化建筑结构和人群
floors, rows, cols = BUILDING_SIZE
building = np.zeros(BUILDING_SIZE, dtype=int)  # 0 表示空，1 表示墙/不可通过
people = []

for f in range(floors):
    for i in range(rows):
        for j in range(cols):
            if random.random() < POPULATION_DENSITY:
                people.append(Person((f, i, j)))

# 模拟主函数
def is_valid(pos):
    f, x, y = pos
    return 0 <= f < floors and 0 <= x < rows and 0 <= y < cols and building[pos] == 0

def get_neighbors(pos):
    f, x, y = pos
    directions = [(0,1,0), (0,-1,0), (0,0,1), (0,0,-1), (1,0,0), (-1,0,0)]
    neighbors = []
    for df, dx, dy in directions:
        new = (f+df, x+dx, y+dy)
        if is_valid(new):
            neighbors.append(new)
    return neighbors

def congestion_map():
    cmap = np.zeros(BUILDING_SIZE)
    for p in people:
        if not p.finished:
            cmap[p.pos] += 1
    return cmap

# 简单启发函数：根据距离和拥堵评分寻找最优路径
from queue import PriorityQueue

def heuristic(p, target):
    return abs(p[0]-target[0]) + abs(p[1]-target[1]) + abs(p[2]-target[2])

def find_path(start, cmap):
    visited = set()
    q = PriorityQueue()
    q.put((0, start, []))

    while not q.empty():
        cost, current, path = q.get()
        if current in visited:
            continue
        visited.add(current)
        path = path + [current]
        if current in EXITS:
            return path

        for neighbor in get_neighbors(current):
            h_cost = heuristic(neighbor, EXITS[0])  # 这里默认用第一个出口启发
            congestion_penalty = GAMMA * cmap[neighbor]
            energy_cost = ALPHA if neighbor[0] == current[0] else BETA
            total_cost = cost + energy_cost + congestion_penalty + h_cost
            q.put((total_cost, neighbor, path))
    return []

# 动画可视化
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
scat = ax.scatter([], [], [], c='blue')

ax.set_xlim(0, rows)
ax.set_ylim(0, cols)
ax.set_zlim(0, floors)

frames = []

def update(frame):
    ax.clear()
    cmap = congestion_map()
    for p in people:
        if p.finished or p.energy <= 0:
            continue
        path = find_path(p.pos, cmap)
        if len(path) > 1:
            next_pos = path[1]
            move_cost = BETA if next_pos[0] != p.pos[0] else ALPHA
            congestion_penalty = GAMMA * cmap[next_pos]
            total_cost = move_cost + congestion_penalty
            p.move(next_pos, total_cost)

    xs = [p.pos[1] for p in people if not p.finished]
    ys = [p.pos[2] for p in people if not p.finished]
    zs = [p.pos[0] for p in people if not p.finished]
    ax.scatter(xs, ys, zs, c='red', marker='o')
    ax.set_title(f"Step {frame} 剩余人数: {len([p for p in people if not p.finished])}")
    ax.set_xlim(0, rows)
    ax.set_ylim(0, cols)
    ax.set_zlim(0, floors)

ani = FuncAnimation(fig, update, frames=200, interval=300)
plt.show()
